import MNN

from torchvision import transforms
from PIL import Image

import MNN.expr as F
import MNN
import cv2
import torch
import numpy as np

# 0-9有10个数字,输出的结果即为识别的数字

mnn_model_path = 'model.mnn'
image_path = '8.jpg'

interpreter = MNN.Interpreter(mnn_model_path)
session = interpreter.createSession()
input_tensor = interpreter.getSessionInput(session)

input_image = cv2.imread(image_path,0)
print("shape:{}".format(input_image.shape))
# cv2.imshow("demo",input_image)
input_image = cv2.resize(input_image,(28,28))
cv2.imshow("img",input_image)
cv2.waitKey(0)

print("shape:{}".format(input_image.shape))
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.1307], [0.3081])
])

tmp_input = preprocess(input_image)

data = tmp_input.numpy()
print("data:{}".format(data.shape))
# np.save("./yuwei.npy", data)

tmp_inputp = MNN.Tensor((1, 1, 28, 28), MNN.Halide_Type_Float,\
                    data, MNN.Tensor_DimensionType_Caffe)

# print("tmp_inputp : {}".format(tmp_inputp.getData()))

input_tensor.copyFrom(tmp_inputp)
# print("3333")
interpreter.runSession(session)
output_tensor = interpreter.getSessionOutput(session)


# print(" output : {}".format(output_tensor.getData()))

print(" output belong to class: {}".format(np.argmax(output_tensor.getData())))